
import string
import random
from sm4 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT
import func


crypt_sm4 = CryptSM4()


def gen_key(bits):
    #key = random.randbytes(15)
    #key = random.sample('zyxwvutsrqponmlkjihgfedcba1234567890', 5)
    num_set = [chr(i) for i in range(48, 58)]
    char_set = [chr(i) for i in range(97, 123)]
    total_set = num_set + char_set
    key = "".join(random.sample(total_set, bits))
    #key = ''.join(random.sample(string.ascii_letters + string.digits, 15))
    return key


def iv_gen(bits):
    num_set = [chr(i) for i in range(48, 58)]
    char_set = [chr(i) for i in range(97, 123)]
    total_set = num_set + char_set
    iv = "".join(random.sample(total_set, bits))
    #iv = ''.join(random.sample(string.ascii_letters + string.digits, 15))
    return iv


def SM4_ecb_enc(plaintext, key):
    crypt_sm4.set_key(key, SM4_ENCRYPT)
    ciphertext_bytes = crypt_sm4.crypt_ecb(plaintext)

    return ciphertext_bytes


def SM4_ecb_dec(ciphertext, key):
    crypt_sm4 = CryptSM4()
    crypt_sm4.set_key(key, SM4_DECRYPT)
    plaintext_bytes = crypt_sm4.crypt_ecb(ciphertext)

    return plaintext_bytes


def SM4_cbc_enc(plaintext, iv, key):
    crypt_sm4.set_key(key, SM4_ENCRYPT)
    ciphertext_bytes = crypt_sm4.crypt_cbc(iv, plaintext)  # bytes类型

    return ciphertext_bytes


def SM4_cbc_dec(ciphertext, iv, key):
    crypt_sm4.set_key(key, SM4_DECRYPT)
    plaintext_bytes = crypt_sm4.crypt_cbc(iv, ciphertext)

    return plaintext_bytes


def exp_SM4_str(plaintext):
    plaintext_bytes = bytes(plaintext, encoding="utf8")
    return plaintext_bytes


def exp_SM4_file(filename):
    with open(filename, "rb") as f:
        plaintext_bytes = f.read()
    f.close()
    return plaintext_bytes

def get_key():
    key = gen_key(16)
    key_bytes = bytes(key, encoding="utf8")
    with open('key_bytes.txt', 'wb') as f:
        f.truncate()
        f.write(key_bytes)
    f.close
    print("key为：", key)
    return key_bytes
def get_iv():
    iv = iv_gen(16)
    iv_bytes = bytes(iv, encoding="utf8")    
    return iv_bytes

    